Skip to content

feat(train): hook in assistant_only_loss for TRL SFT path - #23

Closed
paulpak58 wants to merge 5 commits into
Liquid4All:mainfrom
paulpak58:pp/assistant-only-loss
Closed

feat(train): hook in assistant_only_loss for TRL SFT path#23
paulpak58 wants to merge 5 commits into
Liquid4All:mainfrom
paulpak58:pp/assistant-only-loss

Conversation

@paulpak58

Copy link
Copy Markdown

No description provided.

@paulpak58
paulpak58 marked this pull request as ready for review August 27, 2026 17:43
@paulpak58

Copy link
Copy Markdown
Author

failing lint CI exists on main

Comment thread lqh/tools/handlers.py Outdated
from lqh.models import is_vlm_model_name

is_vision = is_vlm_model_name(base_model)
if assistant_only_loss and (is_vision or type != "sft"):

@leigh-liquid leigh-liquid Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Validate assistant-mask support before accepting the launch. This guard accepts every text model, but the curated and recommended fine-tuning base LiquidAI/LFM2.5-350M-Base ships a chat template with no generation block. With the pinned TRL 1.0 stack, SFTTrainer then raises during dataset tokenization, after the job has already been submitted and provisioned. Please validate or patch the actual tokenizer template, including local checkpoint paths, or reject known-incompatible models before launch.

Comment thread lqh/train/sft.py
gradient_checkpointing=training_cfg.get("gradient_checkpointing", True),
bf16=training_cfg.get("bf16", True),
max_length=training_cfg.get("max_seq_length", 2048),
assistant_only_loss=bool(training_cfg.get("assistant_only_loss", False)),

@leigh-liquid leigh-liquid Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Avoid truncating away every trainable token. SFTConfig defaults to truncation_mode keep_start. With assistant_only_loss enabled and the existing 2048-token limit, a long user turn can leave all assistant tokens beyond the cutoff. TRL checks the mask before truncation, so this is not rejected. I reproduced a 3013-token LFM2.5 row that yielded zero labels other than -100, which can produce zero or NaN-loss batches exactly for the user-heavy rows this option targets. Set truncation_mode to keep_end when this flag is enabled, or explicitly reject or drop rows with no assistant labels after truncation.


@pytest.fixture(scope="module")
def tokenizer():
transformers = pytest.importorskip("transformers")

@leigh-liquid leigh-liquid Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Make the masking regression run in CI. The current workflow installs only the default and dev dependency groups, so transformers, torch, and trl all take the importorskip path. Even a train-extra job would normally miss the model because this fixture requires a pre-warmed local HF cache. The CI-equivalent run reports this test as skipped, leaving the actual feature path untested. Please use a tiny hermetic tokenizer and template fixture, or add a job that installs the train extras and supplies the pinned tokenizer.

@paulpak58
paulpak58 force-pushed the pp/assistant-only-loss branch from 03cf4ad to 2358e58 Compare September 3, 2026 03:12
mlech26l added a commit that referenced this pull request Sep 3, 2026
Text SFT trained at a fixed max_seq_length of 2048 and silently truncated
longer rows. The length is now derived per run from the dataset: the
longest row, rounded up to 1024, capped at 32768. Short datasets keep
their throughput (the calibration probe measures at the configured
length, so a fixed 32k would have collapsed every run to micro-batch 1);
long ones get what they need. Nothing is exposed to the user: no tool
argument, no mention of sequence or context length in the normal flow.

- defaults: MAX_SEQ_LENGTH_CEILING, derived_seq_length(), auto_seq_length
  flag on HParams; DPO/GRPO/vision keep the fixed 2048.
- seq_length (new): submit-time estimate with the base model's
  tokenizer.json via the `tokenizers` wheel (new core dep), character
  fallback when offline (non-ASCII counted one-to-one); streams parquet.
  The backend planner sizes the GPU from the shipped value.
- sft: exact re-measure through the real chat template before the probe,
  capped at the model's max_position_embeddings; rows over the limit are
  dropped with a count (never truncated); counts persisted; all-eval-
  dropped warned. Probe "no_fit" now fails fast with a plain message
  instead of training on the default micro-batch; a skipped probe at
  long context falls back to micro-batch 1.
- calibrate: maybe_autotune_batch_size returns cached|probed|no_fit|skipped.
- Hidden expert override LQH_MAX_SEQ_LENGTH pins the value.
- Skill text, failure messages and the VLM collator error no longer name
  the knob. Unit conftest stubs the Hub tokenizer fetch.

Supersedes PR #22; PR #23 (assistant_only_loss) rebases onto this and
shares tokenized_row_lengths.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015JpgUssR166axwm95srXbg
@paulpak58 paulpak58 closed this Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants